import sys
import io
import os
import json
import logging
import subprocess
import soundfile as sf
import numpy as np
import time
import warnings
from vosk import Model, KaldiRecognizer

# 修复Windows环境下的编码问题
if sys.platform.startswith('win'):
    # Windows环境下强制设置编码
    import codecs
    sys.stdout = codecs.getwriter('utf-8')(sys.stdout.buffer)
    sys.stderr = codecs.getwriter('utf-8')(sys.stderr.buffer)
    # 设置环境变量
    os.environ['PYTHONIOENCODING'] = 'utf-8'
else:
    # 非Windows环境使用原有方式
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8')

def safe_print(message):
    """安全的打印函数，处理编码问题"""
    try:
        print(message, file=sys.stderr)
    except UnicodeEncodeError:
        # 如果UTF-8编码失败，尝试使用系统默认编码
        print(message.encode('utf-8', errors='ignore').decode('utf-8', errors='ignore'), file=sys.stderr)

def download_vosk_model():
    """下载Vosk中文模型"""
    import urllib.request
    import zipfile
    
    model_url = "https://alphacephei.com/vosk/models/vosk-model-small-cn-0.22.zip"
    model_name = "vosk-model-small-cn-0.22"
    
    if os.path.exists(model_name):
        safe_print(f"✅ Vosk中文模型已存在: {model_name}")
        return model_name
    
    safe_print(f"📥 正在下载Vosk中文模型...")
    try:
        # 下载模型
        zip_path = f"{model_name}.zip"
        urllib.request.urlretrieve(model_url, zip_path)
        
        # 解压模型
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(".")
        
        # 删除zip文件
        os.remove(zip_path)
        
        safe_print(f"✅ Vosk中文模型下载完成: {model_name}")
        return model_name
    except Exception as e:
        safe_print(f"❌ 模型下载失败: {e}")
        return None

def load_vosk_model():
    """加载Vosk模型"""
    try:
        # 尝试使用已下载的中文模型
        model_path = download_vosk_model()
        if model_path and os.path.exists(model_path):
            safe_print(f"🔄 正在加载Vosk中文模型: {model_path}")
            model = Model(model_path)
            safe_print("✅ Vosk中文模型加载成功")
            return model
        
        # 如果没有中文模型，尝试使用默认模型
        safe_print("🔄 正在加载Vosk默认模型...")
        model = Model(lang="en-us")  # 使用英文模型作为备用
        safe_print("✅ Vosk默认模型加载成功")
        return model
        
    except Exception as e:
        safe_print(f"❌ Vosk模型加载失败: {e}")
        return None

def recognize_speech_vosk(audio_file_path):
    """使用Vosk进行语音识别"""
    try:
        # 加载模型
        model = load_vosk_model()
        if model is None:
            return {"error": "Vosk模型加载失败"}
        
        # 读取音频文件
        safe_print(f"🎵 正在读取音频文件: {audio_file_path}")
        audio_data, sample_rate = sf.read(audio_file_path)
        
        # 确保音频是单声道
        if len(audio_data.shape) > 1:
            audio_data = audio_data[:, 0]
        
        # 转换为16kHz采样率（Vosk推荐）
        if sample_rate != 16000:
            safe_print(f"🔄 转换采样率从 {sample_rate}Hz 到 16000Hz")
            # 简单的重采样（实际应用中建议使用librosa）
            if sample_rate > 16000:
                # 降采样
                ratio = sample_rate // 16000
                audio_data = audio_data[::ratio]
            else:
                # 升采样（简单重复）
                ratio = 16000 // sample_rate
                audio_data = np.repeat(audio_data, ratio)
            sample_rate = 16000
        
        # 创建识别器
        rec = KaldiRecognizer(model, sample_rate)
        rec.SetWords(True)  # 启用词级别时间戳
        
        # 进行识别
        safe_print("🎤 正在进行语音识别...")
        start_time = time.time()
        
        # 将音频数据转换为字节
        audio_bytes = (audio_data * 32767).astype(np.int16).tobytes()
        
        # 识别
        rec.AcceptWaveform(audio_bytes)
        result = rec.FinalResult()
        
        end_time = time.time()
        processing_time = end_time - start_time
        
        # 解析结果
        try:
            result_json = json.loads(result)
            text = result_json.get("text", "")
            confidence = result_json.get("confidence", 0.0)
            
            safe_print(f"✅ 识别完成！耗时: {processing_time:.2f}秒")
            safe_print(f"📝 识别结果: {text}")
            safe_print(f"🎯 置信度: {confidence:.2f}")
            
            return {
                "success": True,
                "text": text,
                "confidence": confidence,
                "processing_time": processing_time,
                "model": "Vosk",
                "language": "中文" if "vosk-model-small-cn" in str(model) else "英文"
            }
            
        except json.JSONDecodeError:
            safe_print(f"❌ 结果解析失败: {result}")
            return {"error": f"结果解析失败: {result}"}
            
    except Exception as e:
        safe_print(f"❌ Vosk语音识别失败: {e}")
        return {"error": f"Vosk语音识别失败: {e}"}

def main():
    """主函数"""
    if len(sys.argv) != 2:
        safe_print("❌ 用法: python speech_recognize_vosk.py <音频文件路径>")
        sys.exit(1)
    
    audio_file_path = sys.argv[1]
    
    if not os.path.exists(audio_file_path):
        safe_print(f"❌ 音频文件不存在: {audio_file_path}")
        sys.exit(1)
    
    safe_print("🎤 开始Vosk语音识别...")
    safe_print("=" * 50)
    
    # 进行语音识别
    result = recognize_speech_vosk(audio_file_path)
    
    # 输出结果
    if result.get("success"):
        safe_print("=" * 50)
        safe_print("🎉 识别成功！")
        safe_print(f"📝 文本: {result['text']}")
        safe_print(f"🎯 置信度: {result['confidence']:.2f}")
        safe_print(f"⏱️ 处理时间: {result['processing_time']:.2f}秒")
        safe_print(f"🤖 模型: {result['model']}")
        safe_print(f"🌍 语言: {result['language']}")
        
        # 输出JSON格式结果
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        safe_print("=" * 50)
        safe_print("❌ 识别失败！")
        safe_print(f"错误: {result.get('error', '未知错误')}")
        sys.exit(1)

if __name__ == "__main__":
    main()
